#include <assert.h>
#include <ctype.h>
#include <stdio.h>
#include <string.h>
#include "esmd_types.h"
#include "list.h"
#include "memory_cpp.hpp"
#include "io_psf.h"
#define PSF_MAX_LINE_LEN 65536
#define SECTION_NOT_FOUND -1L
//this is only a subset of http://www.charmm-gui.org/?doc=lecture&module=pdb&lesson=6
static long find_section(FILE *fpsf, const char *title) {
  char line[PSF_MAX_LINE_LEN];
  size_t titlelen = strlen(title);
  while (fgets(line, PSF_MAX_LINE_LEN, fpsf)) {
    char *p = line;
    while (isspace(*p))
      p++;
    if (!isdigit(*p))
      continue;
    long count = 0;
    while (isdigit(*p)) {
      count = count * 10 + (unsigned char)(*p - '0');
      p++;
    }
    while (isspace(*p))
      p++;
    if (!strncmp(p, title, titlelen))
      return count;
  }
  return SECTION_NOT_FOUND;
}

psf_data_t *read_psf(FILE *fpsf) {
  char line[PSF_MAX_LINE_LEN];
  psf_data_t *psfdat = esmd::allocate<psf_data_t>(1, "PSF/header");
  psfdat->natom = find_section(fpsf, "!NATOM");
  assert(psfdat->natom != SECTION_NOT_FOUND);
  psfdat->atom = esmd::allocate<psf_atom_t>(psfdat->natom, "PSF/atoms");
  for (long i = 0; i < psfdat->natom; i++) {
    psf_atom_t *atom = psfdat->atom + i;
    fscanf(fpsf, F_GI "%*s%*d%*s%*s%s" F_LR F_LR, &atom->id, atom->type, &atom->charge, &atom->mass);
    fgets(line, PSF_MAX_LINE_LEN, fpsf);
  }

  psfdat->nbond = find_section(fpsf, "!NBOND: bonds");
  assert(psfdat->nbond != SECTION_NOT_FOUND);
  psfdat->bond = esmd::allocate<bond_t>(psfdat->nbond, "PSF/bonds");
  for (long i = 0; i < psfdat->nbond; i++) {
    long *bond = psfdat->bond[i];
    fscanf(fpsf, REP2(F_GI), &bond[0], &bond[1]);
  }
  fgets(line, PSF_MAX_LINE_LEN, fpsf);

  psfdat->nangle = find_section(fpsf, "!NTHETA: angles");
  assert(psfdat->nangle != SECTION_NOT_FOUND);
  psfdat->angle = esmd::allocate<angle_t>(psfdat->nangle, "PSF/angles");
  for (long i = 0; i < psfdat->nangle; i++) {
    long *angle = psfdat->angle[i];
    fscanf(fpsf, REP3(F_GI), &angle[0], &angle[1], &angle[2]);
  }
  fgets(line, PSF_MAX_LINE_LEN, fpsf);

  psfdat->ntori = find_section(fpsf, "!NPHI: dihedrals");
  assert(psfdat->ntori != SECTION_NOT_FOUND);
  psfdat->tori = esmd::allocate<dihed_t>(psfdat->ntori, "PSF/torsions");
  for (long i = 0; i < psfdat->ntori; i++) {
    long *tori = psfdat->tori[i];
    fscanf(fpsf, REP4(F_GI), &tori[0], &tori[1], &tori[2], &tori[3]);
  }
  fgets(line, PSF_MAX_LINE_LEN, fpsf);

  psfdat->nimpr = find_section(fpsf, "!NIMPHI: impropers");
  if (psfdat->nimpr != SECTION_NOT_FOUND) {
  psfdat->impr = esmd::allocate<dihed_t>(psfdat->nimpr, "PSF/impropers");
  for (long i = 0; i < psfdat->nimpr; i++) {
    long *impr = psfdat->impr[i];
    fscanf(fpsf, REP4(F_GI), &impr[0], &impr[1], &impr[2], &impr[3]);
  }
  fgets(line, PSF_MAX_LINE_LEN, fpsf);
  } else {
    psfdat->nimpr = 0;
  }
  // psfdat->ndon = find_section(fpsf, "!NDON: donors");
  // assert(psfdat->ndon != SECTION_NOT_FOUND);
  // if (psfdat->ndon > 0) {
  //   psfdat->don = esmd::allocate<bond_t>(psfdat->ndon, "PSF/donors");
  //   for (long i = 0; i < psfdat->ndon; i++) {
  //     //i: donor, j: hydrogen
  //     long *don = psfdat->don[i];
  //     fscanf(fpsf, REP2(F_GI), &don[0], &don[1]);
  //   }
  // }
  // fgets(line, PSF_MAX_LINE_LEN, fpsf);

  // psfdat->nacc = find_section(fpsf, "!NACC: acceptors");
  // assert(psfdat->nacc != SECTION_NOT_FOUND);
  // if (psfdat->nacc > 0) {
  //   psfdat->acc = esmd::allocate<bond_t>(psfdat->nacc, "PSF/acceptors");
  //   for (long i = 0; i < psfdat->nacc; i++) {
  //     //i: acceptor ancestor, j: acceptor
  //     long *acc = psfdat->acc[i];
  //     fscanf(fpsf, REP2(F_GI), &acc[0], &acc[1]);
  //   }
  // }
  // fgets(line, PSF_MAX_LINE_LEN, fpsf);

  // psfdat->nexcl = find_section(fpsf, "!NNB");
  // if (psfdat->nexcl > 0) {
  //   psfdat->excl = esmd::allocate<bond_t>(psfdat->nexcl, "PSF exclusions");
  //   for (long i = 0; i < psfdat->nexcl; i++) {
  //     long *excl = psfdat->excl[i];
  //     fscanf(fpsf, REP2(F_GI), &excl[0], &excl[1]);
  //   }
  //   fgets(line, PSF_MAX_LINE_LEN, fpsf);
  // }
  return psfdat;
}

void free_psf(psf_data_t *psf) {
  esmd::deallocate(psf->atom);
  esmd::deallocate(psf->bond);
  esmd::deallocate(psf->angle);
  esmd::deallocate(psf->tori);
  esmd::deallocate(psf->impr);
  if (psf->nexcl > 0)
    esmd::deallocate(psf->excl);
  if (psf->nacc > 0)
    esmd::deallocate(psf->acc);
  if (psf->ndon > 0)
    esmd::deallocate(psf->don);
  esmd::deallocate(psf);
}
#ifdef TEST_PSF
int main() {
  memory_init();
  FILE *fpsf = fopen("ubq_wb/ubq_wb.psf", "r");
  psf_data_t *psfdat = read_psf(fpsf);
  for (int i = 0; i < psfdat->natom; i++) {
    printf("ATOM: %ld %s %f %f\n", psfdat->atom[i].id, psfdat->atom[i].type, psfdat->atom[i].charge, psfdat->atom[i].mass);
  }

  for (int i = 0; i < psfdat->nbond; i++) {
    printf("BOND: %ld %ld\n", psfdat->bond[i].i, psfdat->bond[i].j);
  }

  for (int i = 0; i < psfdat->nangle; i++) {
    printf("ANGLE: %ld %ld %ld\n", psfdat->angle[i].i, psfdat->angle[i].j, psfdat->angle[i].k);
  }

  for (int i = 0; i < psfdat->ntori; i++) {
    printf("PHI: %ld %ld %ld %ld\n", psfdat->tori[i].i, psfdat->tori[i].j, psfdat->tori[i].k, psfdat->tori[i].l);
  }

  for (int i = 0; i < psfdat->nimpr; i++) {
    printf("IMPH: %ld %ld %ld %ld\n", psfdat->impr[i].i, psfdat->impr[i].j, psfdat->impr[i].k, psfdat->impr[i].l);
  }
  memory_print();
}
#endif